#!/usr/bin/env python3
"""
Read IMU from a ros2 bag (.db3), convert linear_acceleration to world frame using
IMU orientation, detect static periods, compute simple static bias (world) and
plot x/y/z world accelerations over time with mean and ±2σ lines.

Usage:
  python3 demo/plot_imu_acc_world.py [--bag PATH] [--seconds N]

What this script does (high level):
- Opens the ros2 db3 storage and reads messages from /unilidar/imu
- Converts each linear_acceleration from body frame to world frame using the
  orientation quaternion in the IMU message.
- Records timestamps (seconds since start) and the world-frame acceleration.
- Detects static segments where |acc - gravity| is small and angular velocity
  is small for a minimum duration.
- Estimates a simple static bias from static segments (mean offset from gravity)
  and prints the bias and suggested gravity compensation.
- Plots three vertically stacked matplotlib axes (x, y, z world accelerations
  vs time), each with mean and ±2σ lines and shaded static segments.

Notes and caveats:
- This is an offline analysis/visualization helper, not a real-time estimator.
- The static bias estimate is simplistic (mean over detected static samples).
- For production-grade inertial navigation you'd use calibration, ZUPT, or
  sensor fusion with lidar/vision/odom.
"""

import argparse
import numpy as np
import matplotlib.pyplot as plt
from rosbag2_py import SequentialReader, StorageOptions, ConverterOptions, StorageFilter
from sensor_msgs.msg import Imu
from rclpy.serialization import deserialize_message

G = 9.80665  # gravity magnitude used for checks (m/s^2)
# 重力加速度
GRAVITY = np.array([0.0, 0.0, 9.66723215623791], dtype=np.float32)

# quaternion to rotation matrix (same convention used across demos)
def quat2R(qw, qx, qy, qz):
    """Convert quaternion (w,x,y,z) to a 3x3 rotation matrix."""
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float64)
    return R


def read_imu_world_acc(bag_path, max_seconds=None):
    """Read IMU messages and return times (s since start), world accelerations,
    and raw imu messages (for angular velocity checks).

    Returns:
      times: numpy array (N,)
      acc_w: numpy array (N,3) world-frame accelerations in m/s^2
      ang_v: numpy array (N,3) angular velocities (rad/s)
    """
    reader = SequentialReader()
    reader.open(StorageOptions(uri=bag_path, storage_id='sqlite3'), ConverterOptions('', ''))
    reader.set_filter(StorageFilter(topics=['/unilidar/imu']))

    times = []
    accs = []
    angvs = []
    start_time = None

    while reader.has_next():
        topic, data, t = reader.read_next()
        if topic != '/unilidar/imu':
            continue
        imu = deserialize_message(data, Imu)
        # compute timestamp in seconds (use header if present)
        if hasattr(imu, 'header') and imu.header.stamp.sec != 0:
            sec = imu.header.stamp.sec + imu.header.stamp.nanosec * 1e-9
        else:
            # fallback to rosbag read time
            sec = t * 1e-9
        if start_time is None:
            start_time = sec
        rel_t = sec - start_time
        if max_seconds is not None and rel_t > max_seconds:
            break

        # orientation
        q = imu.orientation
        R = quat2R(q.w, q.x, q.y, q.z)
        # body-frame acceleration
        a_b = np.array([imu.linear_acceleration.x,
                        imu.linear_acceleration.y,
                        imu.linear_acceleration.z], dtype=np.float64)
        # transform to world frame
        a_w = R @ a_b - GRAVITY
        #a_w = R @ a_b --- IGNORE ---
        accs.append(a_w)
        angvs.append([imu.angular_velocity.x, imu.angular_velocity.y, imu.angular_velocity.z])
        times.append(rel_t)

    if len(times) == 0:
        return np.empty((0,)), np.empty((0, 3)), np.empty((0, 3))
    return np.array(times), np.vstack(accs), np.vstack(angvs)


def detect_static_segments(times, accs, angvs, acc_thresh=0.2, ang_thresh=0.05, min_duration=0.5):
    """Detect static segments where acceleration is close to gravity and angular
    velocity is small.

    Strategy:
      - compute |acc - [0,0,G]| (vector norm) and |angv| (norm)
      - mark samples where acc_dev < acc_thresh and angv < ang_thresh
      - group consecutive samples and keep groups longer than min_duration

    Returns list of (start_time, end_time, indices array)
    """
    gravity = np.array([0.0, 0.0, G])
    acc_dev = np.linalg.norm(accs - gravity, axis=1)
    ang_norm = np.linalg.norm(angvs, axis=1)

    is_static = np.logical_and(acc_dev < acc_thresh, ang_norm < ang_thresh)

    segments = []
    N = len(times)
    i = 0
    while i < N:
        if not is_static[i]:
            i += 1
            continue
        j = i
        while j + 1 < N and is_static[j + 1]:
            j += 1
        start_t = times[i]
        end_t = times[j]
        if (end_t - start_t) >= min_duration:
            indices = np.arange(i, j + 1)
            segments.append((start_t, end_t, indices))
        i = j + 1
    return segments


def plot_acc_time(times, accs, static_segments, bias_world=None):
    """Plot x,y,z world accelerations vs time in vertically aligned subplots.

    - Draw mean and ±2σ lines on each subplot
    - Shade each static segment with light gray background
    - If bias_world provided, display it in the title/console
    """
    fig, axes = plt.subplots(3, 1, sharex=True, figsize=(12, 8))
    labels = ['ax (m/s^2)', 'ay (m/s^2)', 'az (m/s^2)']
    colors = ['tab:blue', 'tab:orange', 'tab:green']

    means = np.mean(accs, axis=0)
    stds = np.std(accs, axis=0)

    for i, ax in enumerate(axes):
        ax.plot(times, accs[:, i], color=colors[i], linewidth=0.6)
        mu = means[i]
        sigma = stds[i]
        ax.axhline(mu, color='k', linestyle='--', linewidth=1, label=f'mean={mu:.3f}')
        ax.axhline(mu + 2 * sigma, color='gray', linestyle=':', linewidth=1, label=f'+2σ={mu+2*sigma:.3f}')
        ax.axhline(mu - 2 * sigma, color='gray', linestyle=':', linewidth=1, label=f'-2σ={mu-2*sigma:.3f}')
        ax.set_ylabel(labels[i])
        ax.legend(loc='upper right')

        # shade static segments
        for (s, e, idx) in static_segments:
            ax.axvspan(s, e, color='lightgray', alpha=0.5)

    axes[-1].set_xlabel('Time (s)')
    title = 'World-frame IMU accelerations'
    if bias_world is not None:
        title += f' | estimated bias (world)=({bias_world[0]:.3f},{bias_world[1]:.3f},{bias_world[2]:.3f})'
    fig.suptitle(title)
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    plt.show()


def main():
    parser = argparse.ArgumentParser(description='Plot IMU world accelerations from a ros2 db3 bag')
    parser.add_argument('--bag', default='/gitlab/pcd_analysis/demo_data/unilidar-2023-09-22-12-42-04.db3', help='ros2 bag db3 path')
    parser.add_argument('--seconds', type=float, default=None, help='Only read first N seconds, default all')
    parser.add_argument('--acc-thresh', type=float, default=0.2, help='acc deviation threshold for static detection (m/s^2)')
    parser.add_argument('--ang-thresh', type=float, default=0.05, help='angular velocity threshold for static detection (rad/s)')
    parser.add_argument('--min-duration', type=float, default=0.5, help='minimum static duration (s)')
    args = parser.parse_args()

    print('Reading IMU and computing world-frame accelerations...')
    times, accs, angvs = read_imu_world_acc(args.bag, max_seconds=args.seconds)
    if times.size == 0:
        print('No IMU data found in bag. Exiting.')
        return

    # basic statistics
    means = accs.mean(axis=0)
    stds = accs.std(axis=0)
    print(f'Samples: {len(times)}')
    print(f'Mean world acc: x={means[0]:.4f}, y={means[1]:.4f}, z={means[2]:.4f}')
    print(f'Std  world acc: x={stds[0]:.4f}, y={stds[1]:.4f}, z={stds[2]:.4f}')

    # check z near gravity
    print('\nCheck Z-axis vs gravity:')
    print(f'  mean az: {means[2]:.4f} m/s^2; gravity constant G={G:.5f} m/s^2')
    if abs(means[2] - G) < 1.0:
        print('  az is within 1 m/s^2 of +G -> likely z axis points up in world frame.')
    elif abs(means[2] + G) < 1.0:
        print('  az is within 1 m/s^2 of -G -> likely z axis points down in world frame.')
    else:
        print('  az not close to ±G -> orientation/gravity compensation or sensor bias may be off.')

    # detect static segments
    segments = detect_static_segments(times, accs, angvs,
                                      acc_thresh=args.acc_thresh,
                                      ang_thresh=args.ang_thresh,
                                      min_duration=args.min_duration)

    print(f'Found {len(segments)} static segments (min duration {args.min_duration}s)')

    # gather static samples and compute bias estimate if any
    if len(segments) > 0:
        all_idx = np.concatenate([s[2] for s in segments])
        static_acc = accs[all_idx]
        static_mean = static_acc.mean(axis=0)
        static_std = static_acc.std(axis=0)
        print('\nStatic segment accelerations (world frame):')
        print(f'  mean: x={static_mean[0]:.4f}, y={static_mean[1]:.4f}, z={static_mean[2]:.4f}')
        print(f'  std : x={static_std[0]:.4f}, y={static_std[1]:.4f}, z={static_std[2]:.4f}')

        # simple bias estimate: how far static mean is from gravity vector
        gravity = np.array([0.0, 0.0, G])
        bias_world = static_mean - gravity
        print(f'\nEstimated static bias (world) = static_mean - gravity: ({bias_world[0]:.4f}, {bias_world[1]:.4f}, {bias_world[2]:.4f})')
        print(f'Gravity compensation (to apply to acc_world) = -bias_world - gravity (i.e. remove gravity and bias).')
    else:
        bias_world = None
        print('No static segments identified, cannot compute static bias automatically.')

    # plot accelerations with static shading and bias annotation
    plot_acc_time(times, accs, segments, bias_world=bias_world)


if __name__ == '__main__':
    main()
